#include "structMeshDistribution.hpp"
#include "structMeshPartitioning.hpp"
#include "structMesh_testFunctions.hpp"
#include "mpiHelpers.hpp"
#include <catch2/catch_test_macros.hpp>
#include <cmath>
#include <vector>

using Catch::Approx;
using namespace hms;

auto createTopo() -> std::tuple<MPI_Comm_H, Array2i, Array2i, Array2i, Array2i>
{
	MPI_Comm_H comm { MPI_COMM_WORLD };

	/* this basically does the same as creating an MPI topology */
	Array2i nParts;
	nParts.y() = roundToFactor( sqrt(comm.size()), factors(comm.size()) );
	nParts.x() = comm.size()/nParts.y();

	REQUIRE( nParts.prod() == comm.size() );

	Array2i localSize { 4, 4 };
	Array2i globalSize { localSize * nParts };
	Array2i offset { structMeshInitOffset(globalSize, nParts, comm) };
	return { comm, nParts, globalSize, localSize, offset };
}

template<typename MeshType>
void redist(
	StructMeshBase<MeshType>& mesh,
	Array2i& localSize,
	Array2i& offset,
	const Array2i& nParts,
	const MPI_Comm_H& comm
){
	Array2Xi allSizes ( 2, comm.size() );
	MPI_Allgather(
		localSize.data(), 2, MPI_INDEX,
		allSizes.data(), 2, MPI_INDEX, comm
	);
	Array2Xi allOffsets { meshPartSizesToOffsets(nParts, allSizes) };

	Array2i part { comm.rank() % nParts.x(), comm.rank() / nParts.x() };

	/* we reduce the local mesh size of all parts by 1 and add that to
		* the last ranks */
	localSize -= 1;
	for ( int dim{0}; dim<2; ++dim )
		if ( part(dim) == nParts(dim) - 1 )
			localSize(dim) += nParts(dim);
	
	/* reducing the sizes by 1 means that the new offsets are reduced
		* by 1 for each part that came before */
	offset -= part;

	redistributeMesh(mesh.derived(), comm,
		part, nParts,
		localSize, offset, allSizes, allOffsets
	);
}

TEST_CASE( "(Re-)Distributing RectMesh" ){
	auto [comm, nParts, globalSize, localSize, offset] = createTopo();

	static constexpr Index ng { 1 };
	
	auto test = [&]( Array2s&& diff ){

		std::array<ArrayXs,2> coords;
		for ( int dim{0}; dim<2; ++dim ){
			coords[dim] = ArrayXs::LinSpaced(
				/* due to one vertex overlap between partitions, it's just one
				 * extra vertex per dimension in total for the global mesh */
				globalSize(dim) + 1 + 2*ng,
				-ng*diff(dim),
				(globalSize(dim) + ng)*diff(dim) );
			coords[dim](ng) = 0;
		}
		
		RectMesh mesh;

		if ( comm.isMaster() ){
			ArrayXs xc { coords[0] }, yc { coords[1] };
			mesh = RectMesh{std::move(xc), std::move(yc), ng, true};
		}

		mesh = distributeMesh(
			mesh, comm, globalSize, localSize, offset, ng
		);

		SECTION(
			"Distribution - " + std::to_string(ng) + " ghost layers"
		){
			REQUIRE( ( mesh.xCoords() == coords[0].segment(
				offset.x(), localSize.x() + 1 + 2*ng) ).all() );
			REQUIRE( ( mesh.yCoords() == coords[1].segment(
				offset.y(), localSize.y() + 1 + 2*ng) ).all() );
		}
		SECTION(
			"Redistribution - " + std::to_string(ng) + " ghost layers"
		){
			redist(mesh, localSize, offset, nParts, comm);

			REQUIRE( ( mesh.xCoords() == coords[0].segment(
				offset.x(), localSize.x() + 1 + 2*ng) ).all() );
			REQUIRE( ( mesh.yCoords() == coords[1].segment(
				offset.y(), localSize.y() + 1 + 2*ng) ).all() );
		}
	};

	test( {0.1, 0.2} );
}

TEST_CASE( "(Re-)Distributing Non-uniform RectMesh" ){
	MPI_Comm_H comm {MPI_COMM_WORLD};

	if ( comm.size() != 4 ){
		if ( comm.isMaster() )
			std::cerr << "WARNING: tests relating to true StructMeshes and RectMeshes are only "
			"executed if the number of processes is set to 4. Otherwise they will be skipped.";
	}
	else{
		Array2i globalSize {10, 10};
		Array2i nParts     {2, 2};
		Array2i offset     {structMeshInitOffset(globalSize, nParts, comm)};
		Array2i localSize  {5,5};
		

		auto test = [&](int ng){
			RectMesh mesh;
			if( comm.isMaster()) {
				mesh = createTrueRectMeshForMpi(ng);
			}
			std::array<ArrayXs, 2> coords;
			
			coords[0].resize(globalSize(0)+1+2*ng);
			coords[1].resize(globalSize(1)+1+2*ng);
			
			for(Index i{0}; i<ng; i++) {
				// add ghost cells
				coords[0](i)       = -2 * (ng-i);
				coords[0](ng+11+i) = 24 + (3 * (i+1));
				coords[1](i)       = -2 * (ng-i);
				coords[1](ng+11+i) = 23 + (3 * (i+1));
			}
			coords[0](ng + 0) = 0;
			coords[0](ng + 1) = 2;
			coords[0](ng + 2) = 4;
			coords[0](ng + 3) = 6;
			coords[0](ng + 4) = 10;
			coords[0](ng + 5) = 14;
			coords[0](ng + 6) = 15;
			coords[0](ng + 7) = 16;
			coords[0](ng + 8) = 19;
			coords[0](ng + 9) = 21;
			coords[0](ng +10) = 24;
			
			coords[1](ng + 0) = 0;
			coords[1](ng + 1) = 2;
			coords[1](ng + 2) = 6;
			coords[1](ng + 3) = 8;
			coords[1](ng + 4) = 9;
			coords[1](ng + 5) = 13;
			coords[1](ng + 6) = 15;
			coords[1](ng + 7) = 18;
			coords[1](ng + 8) = 19;
			coords[1](ng + 9) = 20;
			coords[1](ng + 10) = 23;

			mesh = distributeMesh(mesh, comm, globalSize, localSize, offset);
			
			SECTION(
				"Distribution - " + std::to_string(ng) + " ghost layers"
			){
				REQUIRE( ( mesh.xCoords() == coords[0].segment(
					offset.x(), localSize.x() + 1 + 2*ng) ).all() );
				REQUIRE( ( mesh.yCoords() == coords[1].segment(
					offset.y(), localSize.y() + 1 + 2*ng) ).all() );
			}
			SECTION(
				"Redistribution - " + std::to_string(ng) + " ghost layers"
			){
				redist(mesh, localSize, offset, nParts, comm);

				REQUIRE( ( mesh.xCoords() == coords[0].segment(
					offset.x(), localSize.x() + 1 + 2*ng) ).all() );
				REQUIRE( ( mesh.yCoords() == coords[1].segment(
					offset.y(), localSize.y() + 1 + 2*ng) ).all() );
			}
		};

		test(1);
	}
}

TEST_CASE( "(Re-)Distributing StructMesh" ){
	auto [comm, nParts, globalSize, localSize, offset] = createTopo();

	auto test = [&]( int ng, Array2s diff ){

		Matrix2Xs vertices ( 2, (globalSize + 1).prod() );
		for (int j{0}; j<=globalSize.y(); ++j){
			for (int i{0}; i<=globalSize.x(); ++i){
				vertices.row(0).reshaped(globalSize.x()+1, globalSize.y()+1)(i, j)
					= i*diff.x();
				vertices.row(1).reshaped(globalSize.x()+1, globalSize.y()+1)(i, j)
					= j*diff.y();
			}
		}

		/* create mesh on all nodes, so that they all have the data to compare
		 * the distributed mesh to */
		StructMesh mesh { std::move(vertices), globalSize };
		
		StructMesh distMesh {
			distributeMesh( mesh, comm, globalSize, localSize, offset )
		};

		auto compareBlocks = [&]( const auto& matA, const auto& matB,
			Index sizeAx, Index sizeAy, Index sizeBx, Index sizeBy,
			Index nVals = 1
		){
			REQUIRE(
				matA.reshaped( nVals*sizeAx, sizeAy ) == 
				matB.reshaped( nVals*sizeBx, sizeBy )
					.block( nVals*offset.x(), offset.y(), nVals*sizeAx, sizeAy )
			);
		};

		auto testMesh = [&](){
			compareBlocks( distMesh.vertices().row(0), mesh.vertices().row(0),
				localSize.x()+1, localSize.y()+1,
				globalSize.x()+1, globalSize.y()+1 );
			compareBlocks( distMesh.vertices().row(1), mesh.vertices().row(1),
				localSize.x() + 1, localSize.y() + 1,
				globalSize.x() + 1, globalSize.y() + 1 );
			/* Edges: The first nx*(ny+1) entries (head/leftCols) are horizontal
			 * edges, the remaining (nx+1)*ny entries (tail/rightCols) are
			 * vertical edges.
			 * Lengths: */
			for ( auto length : distMesh.edges().lengths().head(
					localSize.x() * (localSize.y() + 1) ))
				REQUIRE( length == Approx( diff.x() ) );
			for ( auto length : distMesh.edges().lengths().tail(
					(localSize.x() + 1) * localSize.y() ))
				REQUIRE( length == Approx( diff.y() ) );
			/* Normals: */
			compareBlocks(
				distMesh.edges().normals().leftCols(
					localSize.x() * (localSize.y() + 1) ),
				mesh.edges().normals().leftCols(
					globalSize.x() * (globalSize.y() + 1) ),
				localSize.x(), localSize.y()+1,
				globalSize.x(), globalSize.y() + 1, 2
			);
			compareBlocks( distMesh.edges().normals().rightCols(
					(localSize.x() + 1) * localSize.y() ),
				mesh.edges().normals().rightCols(
					(globalSize.x() + 1) * globalSize.y() ),
				localSize.x() + 1, localSize.y(),
				globalSize.x() + 1, globalSize.y(), 2
			);
			/* cells */
			/* because there are no ghost cells at the corners of the mesh,
			 * all arrays relating to cells have uninitialised values at those
			 * positions. When the mesh is created on all ranks as above, but
			 * then distributed from the master rank, those uninitialised
			 * garbage values are distributed as well. That is not an issue by
			 * itself, but will cause differing values on the other nodes.
			 * So instead of testing the whole array, two blocks are tested,
			 * which overlap in the domain area, and cover all ghost cells once.
			 * */
			auto gAreas = mesh.cells().areas().reshaped(
				2*ng + globalSize.x(), 2*ng + globalSize.y() );
			auto lAreas = distMesh.cells().areas().reshaped(
				localSize.x()+2*ng, localSize.y()+2*ng );

			REQUIRE( gAreas.block(
				offset.x()+ng, offset.y(), localSize.x(), localSize.y()+2*ng
			).isApprox( lAreas.middleRows( ng, localSize.x() )) );
			REQUIRE( gAreas.block(
				offset.x(), offset.y()+ng, localSize.x()+2*ng, localSize.y()
			).isApprox(lAreas.middleCols( ng, localSize.y() )) );

			auto gCentroids = mesh.cells().centroids().reshaped(
				2*(globalSize.x() + 2*ng), 2*ng + globalSize.y() );
			auto lCentroids = distMesh.cells().centroids().reshaped(
				2*(localSize.x()+2*ng), localSize.y()+2*ng );
			
			REQUIRE( gCentroids.block(
				2*offset.x()+2*ng, offset.y(), 2*localSize.x(), localSize.y()+2*ng
			).isApprox( lCentroids.middleRows( 2*ng, 2*localSize.x() ) ) );
			REQUIRE( gCentroids.block(
				2*offset.x(), offset.y()+ng, 2*(localSize.x()+2*ng), localSize.y()
			).isApprox( lCentroids.middleCols( ng, localSize.y() ) ) );
		};

		SECTION( "Distribution - " + std::to_string(ng) + " ghost layers" ){
			testMesh();
		}
		SECTION( "Redistribution - " + std::to_string(ng) + " ghost layers" ){
			redist( distMesh, localSize, offset, nParts, comm );
			testMesh();
		}
	};

	test(1, {0.1, 0.2});
}

TEST_CASE( "Distributing Non-rectangular StructMesh" ){
	MPI_Comm_H comm {MPI_COMM_WORLD};

	if ( comm.size() != 4 ){
		if ( comm.isMaster() )
			std::cerr << "WARNING: tests relating to true StructMeshes and RectMeshes are only "
			"executed if the number of processes is set to 4. Otherwise they will be skipped.";
	}
	else{
		auto test =[&](int ng){

			Array2i globalSize {6, 7};
			Array2i nParts     {2, 2};
			Array2i offset     {structMeshInitOffset(globalSize, nParts, comm)};
			Array2i localSize;

			StructMesh mesh {createTrueStructMeshForMpi(ng)};
			StructMesh distMesh;
			for(Index i{0};  i<4;  i++){
			
				if(comm.rank() == 0 || comm.rank() == 1) localSize = {3,3};
				else localSize = {3,4};
			
			}
			distMesh = distributeMesh(mesh, comm, globalSize, localSize, offset, ng);

			auto compareBlocks = [&]( const auto& matA, const auto& matB,
				Index sizeAx, Index sizeAy, Index sizeBx, Index sizeBy,
				Index nVals = 1
			){
				REQUIRE(
					matA.reshaped( nVals*sizeAx, sizeAy ) == 
					matB.reshaped( nVals*sizeBx, sizeBy )
						.block( nVals*offset.x(), offset.y(), nVals*sizeAx, sizeAy )
				);
			};

			auto testMesh = [&](){
				compareBlocks( distMesh.vertices().row(0), mesh.vertices().row(0),
					localSize.x()+1, localSize.y()+1,
					globalSize.x()+1, globalSize.y()+1 );
				compareBlocks( distMesh.vertices().row(1), mesh.vertices().row(1),
					localSize.x() + 1, localSize.y() + 1,
					globalSize.x() + 1, globalSize.y() + 1 );

				// edge lengths
				// if( redistributed == false){
					if( comm.rank() == 0) {
						REQUIRE( distMesh.edge(0 ).length() == Approx(3.0));
						REQUIRE( distMesh.edge(1 ).length() == Approx(2.0));
						REQUIRE( distMesh.edge(2 ).length() == Approx(2.0));
						REQUIRE( distMesh.edge(3 ).length() == Approx(2.0));
						REQUIRE( distMesh.edge(4 ).length() == Approx(1.0));
						REQUIRE( distMesh.edge(5 ).length() == Approx(2.0));
						REQUIRE( distMesh.edge(6 ).length() == Approx(1.0));
						REQUIRE( distMesh.edge(7 ).length() == Approx(sqrt(26)/2));
						REQUIRE( distMesh.edge(8 ).length() == Approx(sqrt(10)/2));
						REQUIRE( distMesh.edge(9 ).length() == Approx(sqrt(5)));
						REQUIRE( distMesh.edge(10).length() == Approx(2));
						REQUIRE( distMesh.edge(11).length() == Approx(2));

						REQUIRE( distMesh.edge(12).length() == Approx(sqrt(5)));
						REQUIRE( distMesh.edge(13).length() == Approx(2));
						REQUIRE( distMesh.edge(14).length() == Approx(sqrt(5)));
						REQUIRE( distMesh.edge(15).length() == Approx(sqrt(5)));
						REQUIRE( distMesh.edge(16).length() == Approx(2));
						REQUIRE( distMesh.edge(17).length() == Approx(sqrt(5)));
						REQUIRE( distMesh.edge(18).length() == Approx(sqrt(10)/2));
						REQUIRE( distMesh.edge(19).length() == Approx(2));
						REQUIRE( distMesh.edge(20).length() == Approx(sqrt(5)));
						REQUIRE( distMesh.edge(21).length() == Approx(1));
						REQUIRE( distMesh.edge(22).length() == Approx(sqrt(10)/2));
						REQUIRE( distMesh.edge(23).length() == Approx(1));
					}

					if(comm.rank() == 1) {
						REQUIRE( distMesh.edge(0 ).length() == Approx(1));
						REQUIRE( distMesh.edge(1 ).length() == Approx(3));
						REQUIRE( distMesh.edge(2 ).length() == Approx(2));
						REQUIRE( distMesh.edge(3 ).length() == Approx(sqrt(5)));
						REQUIRE( distMesh.edge(4 ).length() == Approx(sqrt(5)));
						REQUIRE( distMesh.edge(5 ).length() == Approx(sqrt(10)));
						REQUIRE( distMesh.edge(6 ).length() == Approx(sqrt(5)));
						REQUIRE( distMesh.edge(7 ).length() == Approx(2));
						REQUIRE( distMesh.edge(8 ).length() == Approx(sqrt(13)));
						REQUIRE( distMesh.edge(9 ).length() == Approx(sqrt(5)));
						REQUIRE( distMesh.edge(10).length() == Approx(sqrt(5)));
						REQUIRE( distMesh.edge(11).length() == Approx(2*sqrt(2)));

						REQUIRE( distMesh.edge(12).length() == Approx(sqrt(5)));
						REQUIRE( distMesh.edge(13).length() == Approx(1));
						REQUIRE( distMesh.edge(14).length() == Approx(sqrt(5)));
						REQUIRE( distMesh.edge(15).length() == Approx(3));
						REQUIRE( distMesh.edge(16).length() == Approx(2));
						REQUIRE( distMesh.edge(17).length() == Approx(2));
						REQUIRE( distMesh.edge(18).length() == Approx(1));
						REQUIRE( distMesh.edge(19).length() == Approx(2));
						REQUIRE( distMesh.edge(20).length() == Approx(1));
						REQUIRE( distMesh.edge(21).length() == Approx(3));
						REQUIRE( distMesh.edge(22).length() == Approx(2));
						REQUIRE( distMesh.edge(23).length() == Approx(sqrt(5)));
					}

					if(comm.rank() == 2) {
						REQUIRE( distMesh.edge(0 ).length() == Approx(sqrt(5)));
						REQUIRE( distMesh.edge(1 ).length() == Approx(2));
						REQUIRE( distMesh.edge(2 ).length() == Approx(2));
						REQUIRE( distMesh.edge(3 ).length() == Approx(2));
						REQUIRE( distMesh.edge(4 ).length() == Approx(3));
						REQUIRE( distMesh.edge(5 ).length() == Approx(1));
						REQUIRE( distMesh.edge(6 ).length() == Approx(sqrt(5)));
						REQUIRE( distMesh.edge(7 ).length() == Approx(sqrt(5)));
						REQUIRE( distMesh.edge(8 ).length() == Approx(sqrt(10)));
						REQUIRE( distMesh.edge(9 ).length() == Approx(sqrt(5)));
						REQUIRE( distMesh.edge(10).length() == Approx(sqrt(10)));
						REQUIRE( distMesh.edge(11).length() == Approx(sqrt(5)));
						REQUIRE( distMesh.edge(12).length() == Approx(2));
						REQUIRE( distMesh.edge(13).length() == Approx(3));
						REQUIRE( distMesh.edge(14).length() == Approx(2));

						REQUIRE( distMesh.edge(15).length() == Approx(2));
						REQUIRE( distMesh.edge(16).length() == Approx(3));
						REQUIRE( distMesh.edge(17).length() == Approx(sqrt(10)));
						REQUIRE( distMesh.edge(18).length() == Approx(3));
						REQUIRE( distMesh.edge(19).length() == Approx(sqrt(5)));
						REQUIRE( distMesh.edge(20).length() == Approx(sqrt(10)));
						REQUIRE( distMesh.edge(21).length() == Approx(2*sqrt(5)));
						REQUIRE( distMesh.edge(22).length() == Approx(3));
						REQUIRE( distMesh.edge(23).length() == Approx(5));
						REQUIRE( distMesh.edge(24).length() == Approx(3));
						REQUIRE( distMesh.edge(25).length() == Approx(sqrt(10)));
						REQUIRE( distMesh.edge(26).length() == Approx(5));
						REQUIRE( distMesh.edge(27).length() == Approx(4));
						REQUIRE( distMesh.edge(28).length() == Approx(5));
						REQUIRE( distMesh.edge(29).length() == Approx(4));
						REQUIRE( distMesh.edge(30).length() == Approx(3));
					}

					if(comm.rank() == 3) {
						REQUIRE( distMesh.edge(0 ).length() == Approx(sqrt(5)));
						REQUIRE( distMesh.edge(1 ).length() == Approx(sqrt(5)));
						REQUIRE( distMesh.edge(2 ).length() == Approx(2*sqrt(2)));
						REQUIRE( distMesh.edge(3 ).length() == Approx(1));
						REQUIRE( distMesh.edge(4 ).length() == Approx(3));
						REQUIRE( distMesh.edge(5 ).length() == Approx(2));
						REQUIRE( distMesh.edge(6 ).length() == Approx(1));
						REQUIRE( distMesh.edge(7 ).length() == Approx(sqrt(17)/2));
						REQUIRE( distMesh.edge(8 ).length() == Approx(sqrt(101)/2));
						REQUIRE( distMesh.edge(9 ).length() == Approx(3));
						REQUIRE( distMesh.edge(10).length() == Approx(sqrt(17)/2));
						REQUIRE( distMesh.edge(11).length() == Approx(sqrt(65)/2));
						REQUIRE( distMesh.edge(12).length() == Approx(3));
						REQUIRE( distMesh.edge(13).length() == Approx(2));
						REQUIRE( distMesh.edge(14).length() == Approx(4));

						REQUIRE( distMesh.edge(15).length() == Approx(3));
						REQUIRE( distMesh.edge(16).length() == Approx(sqrt(5)));
						REQUIRE( distMesh.edge(17).length() == Approx(3));
						REQUIRE( distMesh.edge(18).length() == Approx(1));
						REQUIRE( distMesh.edge(19).length() == Approx(3));
						REQUIRE( distMesh.edge(20).length() == Approx(3));
						REQUIRE( distMesh.edge(21).length() == Approx(sqrt(29)/2));
						REQUIRE( distMesh.edge(22).length() == Approx(sqrt(13)));
						REQUIRE( distMesh.edge(23).length() == Approx(5));
						REQUIRE( distMesh.edge(24).length() == Approx(sqrt(29)));
						REQUIRE( distMesh.edge(25).length() == Approx(sqrt(29)));
						REQUIRE( distMesh.edge(26).length() == Approx(sqrt(26)));
						REQUIRE( distMesh.edge(27).length() == Approx(3));
						REQUIRE( distMesh.edge(28).length() == Approx(3));
						REQUIRE( distMesh.edge(29).length() == Approx(3.5));
						REQUIRE( distMesh.edge(30).length() == Approx(3));
					}
										
				// normals
				compareBlocks(
					distMesh.edges().normals().leftCols(
						localSize.x() * (localSize.y() + 1) ),
					mesh.edges().normals().leftCols(
						globalSize.x() * (globalSize.y() + 1) ),
					localSize.x(), localSize.y()+1,
					globalSize.x(), globalSize.y() + 1, 2
				);
				compareBlocks( distMesh.edges().normals().rightCols(
						(localSize.x() + 1) * localSize.y() ),
					mesh.edges().normals().rightCols(
						(globalSize.x() + 1) * globalSize.y() ),
					localSize.x() + 1, localSize.y(),
					globalSize.x() + 1, globalSize.y(), 2
				);
				// cells
				auto gAreas = mesh.cells().areas().reshaped(
					2*ng + globalSize.x(), 2*ng + globalSize.y() );
				auto lAreas = distMesh.cells().areas().reshaped(
					localSize.x()+2*ng, localSize.y()+2*ng );

				REQUIRE( (gAreas.block(
					offset.x()+ng, offset.y(), localSize.x(), localSize.y()+2*ng
				) == lAreas.middleRows( ng, localSize.x() )).all() );
				REQUIRE( (gAreas.block(
					offset.x(), offset.y()+ng, localSize.x()+2*ng, localSize.y()
				) == lAreas.middleCols( ng, localSize.y() )).all() );
				//centroids
				auto gCentroids = mesh.cells().centroids().reshaped(
					2*(globalSize.x() + 2*ng), 2*ng + globalSize.y() );
				auto lCentroids = distMesh.cells().centroids().reshaped(
					2*(localSize.x()+2*ng), localSize.y()+2*ng );

				REQUIRE( gCentroids.block(
					2*offset.x()+2*ng, offset.y(), 2*localSize.x(), localSize.y()+2*ng
				).isApprox( lCentroids.middleRows( 2*ng, 2*localSize.x() ) ) );
				REQUIRE( gCentroids.block(
					2*offset.x(), offset.y()+ng, 2*(localSize.x()+2*ng), localSize.y()
				) == lCentroids.middleCols( ng, localSize.y() ) );
			};

			SECTION( "Distribution - " + std::to_string(ng) + " ghost layers" ){
				testMesh();
				
			}
		};

		test(1);
	}
}

TEST_CASE( "(Re-)Distributing Non-rectangular StructMesh" ){
	MPI_Comm_H comm {MPI_COMM_WORLD};
	if ( comm.size() != 4 ){
		if ( comm.isMaster() )
			std::cerr << "WARNING: tests relating to true StructMeshes and RectMeshes are only "
			"executed if the number of processes is set to 4. Otherwise they will be skipped.";
	}
	else{
  
		auto test =[&](int ng){
			
			Array2i globalSize {6, 7};
			Array2i nParts     {2, 2};
			Array2i offset     {structMeshInitOffset(globalSize, nParts, comm)};
			Array2i localSize;
			StructMesh mesh {createTrueStructMeshForMpi(ng)};
			StructMesh distMesh;
			for(Index i{0};  i<4;  i++){
			
				if(comm.rank() == 0 || comm.rank() == 1) localSize = {3,3};
				else localSize = {3,4};
			
			}
			distMesh = distributeMesh(mesh, comm, globalSize, localSize, offset, ng);
			


			redist(distMesh, localSize, offset, nParts, comm);

			auto compareBlocks = [&]( const auto& matA, const auto& matB,
				Index sizeAx, Index sizeAy, Index sizeBx, Index sizeBy,
				Index nVals = 1
			){
				REQUIRE(
					matA.reshaped( nVals*sizeAx, sizeAy ) == 
					matB.reshaped( nVals*sizeBx, sizeBy )
						.block( nVals*offset.x(), offset.y(), nVals*sizeAx, sizeAy )
				);
			};

			auto testMesh = [&](){
				compareBlocks( distMesh.vertices().row(0), mesh.vertices().row(0),
					localSize.x()+1, localSize.y()+1,
					globalSize.x()+1, globalSize.y()+1 );
				compareBlocks( distMesh.vertices().row(1), mesh.vertices().row(1),
					localSize.x() + 1, localSize.y() + 1,
					globalSize.x() + 1, globalSize.y() + 1 );

				// edge lengths
				if( comm.rank() == 0) {
					REQUIRE( distMesh.edge(0 ).length() == Approx(3.0));
					REQUIRE( distMesh.edge(1 ).length() == Approx(2.0));
					REQUIRE( distMesh.edge(2 ).length() == Approx(2.0));
					REQUIRE( distMesh.edge(3 ).length() == Approx(1.0));
					REQUIRE( distMesh.edge(4 ).length() == Approx(1.0));
					REQUIRE( distMesh.edge(5 ).length() == Approx(sqrt(26)/2));
					REQUIRE( distMesh.edge(6 ).length() == Approx(sqrt(5)));
					REQUIRE( distMesh.edge(7 ).length() == Approx(2));
					REQUIRE( distMesh.edge(8 ).length() == Approx(sqrt(5)));
					REQUIRE( distMesh.edge(9 ).length() == Approx(2));
					REQUIRE( distMesh.edge(10).length() == Approx(sqrt(5)));
					REQUIRE( distMesh.edge(11).length() == Approx(sqrt(10)/2));
				}

				if(comm.rank() == 1) {
					REQUIRE( distMesh.edge(0 ).length() == Approx(2.0));
					REQUIRE( distMesh.edge(1 ).length() == Approx(1));
					REQUIRE( distMesh.edge(2 ).length() == Approx(3));
					REQUIRE( distMesh.edge(3 ).length() == Approx(2));
					REQUIRE( distMesh.edge(4 ).length() == Approx(2.0));
					REQUIRE( distMesh.edge(5 ).length() == Approx(sqrt(5)));
					REQUIRE( distMesh.edge(6 ).length() == Approx(sqrt(5)));
					REQUIRE( distMesh.edge(7 ).length() == Approx(sqrt(10)));
					REQUIRE( distMesh.edge(8 ).length() == Approx(sqrt(10)/2));
					REQUIRE( distMesh.edge(9 ).length() == Approx(sqrt(5)));
					REQUIRE( distMesh.edge(10).length() == Approx(2));
					REQUIRE( distMesh.edge(11).length() == Approx(sqrt(13)));
					REQUIRE( distMesh.edge(12).length() == Approx(sqrt(5)));
					REQUIRE( distMesh.edge(13).length() == Approx(sqrt(5)));
					REQUIRE( distMesh.edge(14).length() == Approx(1));
					REQUIRE( distMesh.edge(15).length() == Approx(sqrt(5)));
					REQUIRE( distMesh.edge(16).length() == Approx(3));
					REQUIRE( distMesh.edge(17).length() == Approx(sqrt(10)/2));
					REQUIRE( distMesh.edge(18).length() == Approx(2));
					REQUIRE( distMesh.edge(19).length() == Approx(2));
					REQUIRE( distMesh.edge(20).length() == Approx(1));
					REQUIRE( distMesh.edge(21).length() == Approx(2));
				}

				if(comm.rank() == 2) {
					REQUIRE( distMesh.edge(0 ).length() == Approx(1.0));
					REQUIRE( distMesh.edge(1 ).length() == Approx(sqrt(26)/2));
					REQUIRE( distMesh.edge(2 ).length() == Approx(sqrt(5)));
					REQUIRE( distMesh.edge(3 ).length() == Approx(2));
					REQUIRE( distMesh.edge(4 ).length() == Approx(2));
					REQUIRE( distMesh.edge(5 ).length() == Approx(3));
					REQUIRE( distMesh.edge(6 ).length() == Approx(sqrt(5)));
					REQUIRE( distMesh.edge(7 ).length() == Approx(sqrt(5)));
					REQUIRE( distMesh.edge(8 ).length() == Approx(sqrt(5)));
					REQUIRE( distMesh.edge(9 ).length() == Approx(sqrt(10)));
					REQUIRE( distMesh.edge(10).length() == Approx(2));
					REQUIRE( distMesh.edge(11).length() == Approx(3));
					REQUIRE( distMesh.edge(12).length() == Approx(sqrt(5)));
					REQUIRE( distMesh.edge(13).length() == Approx(1));
					REQUIRE( distMesh.edge(14).length() == Approx(sqrt(10)/2));
					REQUIRE( distMesh.edge(15).length() == Approx(2));
					REQUIRE( distMesh.edge(16).length() == Approx(3));
					REQUIRE( distMesh.edge(17).length() == Approx(sqrt(10)));
					REQUIRE( distMesh.edge(18).length() == Approx(sqrt(5)));
					REQUIRE( distMesh.edge(19).length() == Approx(sqrt(10)));
					REQUIRE( distMesh.edge(20).length() == Approx(2*sqrt(5)));
					REQUIRE( distMesh.edge(21).length() == Approx(5));
					REQUIRE( distMesh.edge(22).length() == Approx(3));
					REQUIRE( distMesh.edge(23).length() == Approx(sqrt(10)));
					REQUIRE( distMesh.edge(24).length() == Approx(4));
					REQUIRE( distMesh.edge(25).length() == Approx(5));
					REQUIRE( distMesh.edge(26).length() == Approx(4));
				}

				if(comm.rank() == 3) {
					REQUIRE( distMesh.edge(0 ).length() == Approx(sqrt(10)/2) );
					REQUIRE( distMesh.edge(1 ).length() == Approx(sqrt(5)) );
					REQUIRE( distMesh.edge(2 ).length() == Approx(2) );
					REQUIRE( distMesh.edge(3 ).length() == Approx(sqrt(13)) );
					REQUIRE( distMesh.edge(4 ).length() == Approx(2) );
					REQUIRE( distMesh.edge(5 ).length() == Approx(sqrt(5)) );
					REQUIRE( distMesh.edge(6 ).length() == Approx(sqrt(5)) );
					REQUIRE( distMesh.edge(7 ).length() == Approx(2*sqrt(2)) );
					REQUIRE( distMesh.edge(8 ).length() == Approx(1) );
					REQUIRE( distMesh.edge(9 ).length() == Approx(1) );
					REQUIRE( distMesh.edge(10).length() == Approx(3) );
					REQUIRE( distMesh.edge(11).length() == Approx(2) );
					REQUIRE( distMesh.edge(12).length() == Approx(sqrt(10)) );
					REQUIRE( distMesh.edge(13).length() == Approx(1) );
					REQUIRE( distMesh.edge(14).length() == Approx(sqrt(17)/2) );
					REQUIRE( distMesh.edge(15).length() == Approx(sqrt(101)/2) );
					REQUIRE( distMesh.edge(16).length() == Approx(sqrt(5)) );
					REQUIRE( distMesh.edge(17).length() == Approx(3) );
					REQUIRE( distMesh.edge(18).length() == Approx(sqrt(17)/2) );
					REQUIRE( distMesh.edge(19).length() == Approx(sqrt(65)/2) );
					REQUIRE( distMesh.edge(20).length() == Approx(2) );
					REQUIRE( distMesh.edge(21).length() == Approx(3) );
					REQUIRE( distMesh.edge(22).length() == Approx(2) );
					REQUIRE( distMesh.edge(23).length() == Approx(4) );
					REQUIRE( distMesh.edge(24).length() == Approx(sqrt(10)/2) );
					REQUIRE( distMesh.edge(25).length() == Approx(1) );
					REQUIRE( distMesh.edge(26).length() == Approx(3) );
					REQUIRE( distMesh.edge(27).length() == Approx(2) );
					REQUIRE( distMesh.edge(28).length() == Approx(sqrt(5)) );
					REQUIRE( distMesh.edge(29).length() == Approx(sqrt(10)) );
					REQUIRE( distMesh.edge(30).length() == Approx(3) );
					REQUIRE( distMesh.edge(31).length() == Approx(sqrt(5)) );
					REQUIRE( distMesh.edge(32).length() == Approx(3) );
					REQUIRE( distMesh.edge(33).length() == Approx(1) );
					REQUIRE( distMesh.edge(34).length() == Approx(2*sqrt(5)) );
					REQUIRE( distMesh.edge(35).length() == Approx(3) );
					REQUIRE( distMesh.edge(36).length() == Approx(3) );
					REQUIRE( distMesh.edge(37).length() == Approx(sqrt(29)/2) );
					REQUIRE( distMesh.edge(38).length() == Approx(sqrt(13)) );
					REQUIRE( distMesh.edge(39).length() == Approx(sqrt(10)) );
					REQUIRE( distMesh.edge(40).length() == Approx(5) );
					REQUIRE( distMesh.edge(41).length() == Approx(sqrt(29)) );
					REQUIRE( distMesh.edge(42).length() == Approx(sqrt(29)) );
					REQUIRE( distMesh.edge(43).length() == Approx(sqrt(26)) );
					REQUIRE( distMesh.edge(44).length() == Approx(4) );
					REQUIRE( distMesh.edge(45).length() == Approx(3) );
					REQUIRE( distMesh.edge(46).length() == Approx(3) );
					REQUIRE( distMesh.edge(47).length() == Approx(3.5) );
					REQUIRE( distMesh.edge(48).length() == Approx(3) );
				}
				
				// normals
				compareBlocks(
					distMesh.edges().normals().leftCols(
						localSize.x() * (localSize.y() + 1) ),
					mesh.edges().normals().leftCols(
						globalSize.x() * (globalSize.y() + 1) ),
					localSize.x(), localSize.y()+1,
					globalSize.x(), globalSize.y() + 1, 2
				);
				compareBlocks( distMesh.edges().normals().rightCols(
						(localSize.x() + 1) * localSize.y() ),
					mesh.edges().normals().rightCols(
						(globalSize.x() + 1) * globalSize.y() ),
					localSize.x() + 1, localSize.y(),
					globalSize.x() + 1, globalSize.y(), 2
				);
				// cells
				auto gAreas = mesh.cells().areas().reshaped(
					2*ng + globalSize.x(), 2*ng + globalSize.y() );
				auto lAreas = distMesh.cells().areas().reshaped(
					localSize.x()+2*ng, localSize.y()+2*ng );
				if(comm.rank() == 1 || comm.rank() == 2 || comm.rank() == 3) {
					std::cout << lAreas << std::endl;
				}
				
				REQUIRE( (gAreas.block(
					offset.x()+ng, offset.y(), localSize.x(), localSize.y()+2*ng
				) == lAreas.middleRows( ng, localSize.x() )).all() );
				REQUIRE( (gAreas.block(
					offset.x(), offset.y()+ng, localSize.x()+2*ng, localSize.y()
				) == lAreas.middleCols( ng, localSize.y() )).all() );
				//centroids
				auto gCentroids = mesh.cells().centroids().reshaped(
					2*(globalSize.x() + 2*ng), 2*ng + globalSize.y() );
				auto lCentroids = distMesh.cells().centroids().reshaped(
					2*(localSize.x()+2*ng), localSize.y()+2*ng );

				REQUIRE( gCentroids.block(
					2*offset.x()+2*ng, offset.y(), 2*localSize.x(), localSize.y()+2*ng
				).isApprox( lCentroids.middleRows( 2*ng, 2*localSize.x() ) ) );
				REQUIRE( gCentroids.block(
					2*offset.x(), offset.y()+ng, 2*(localSize.x()+2*ng), localSize.y()
				) == lCentroids.middleCols( ng, localSize.y() ) );
				
			};

			SECTION( "Redistribution - " + std::to_string(ng) + " ghost layers" ){
				testMesh();
			}
		};

		test(1);
		}
}
